import pandas as pd
import numpy as np
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.model_selection import train_test_split
import pickle
from pathlib import Path
from typing import Dict, Any, List
import json
import logging
from datetime import datetime

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("SalaryPredictor")

class SalaryPredictor:
    """智能薪酬预测系统"""
    def __init__(self, data_path: str = "b/data/salary_data"):
        self.data_path = Path(data_path)
        self.model_path = self.data_path / "salary_model.pkl"
        self.model = None
        self.features = [
            'years_experience', 
            'education_level',
            'skill_count',
            'industry',
            'location_tier'
        ]
        self._prepare_data()
        
    def _prepare_data(self):
        """准备数据目录和示例数据"""
        self.data_path.mkdir(parents=True, exist_ok=True)
        
        # 如果不存在示例数据则创建
        sample_file = self.data_path / "sample_salaries.csv"
        if not sample_file.exists():
            sample_data = {
                'years_experience': [1, 3, 5, 7, 10],
                'education_level': [2, 3, 4, 4, 5],  # 1-5表示学历等级
                'skill_count': [3, 5, 7, 8, 10],
                'industry': ['IT', 'Finance', 'IT', 'Healthcare', 'Finance'],
                'location_tier': [1, 2, 1, 3, 1],  # 城市等级
                'salary': [8000, 15000, 25000, 30000, 45000]  # 月薪
            }
            pd.DataFrame(sample_data).to_csv(sample_file, index=False)
            logger.info("已创建示例薪酬数据")
    
    def train_model(self):
        """训练薪酬预测模型"""
        try:
            # 加载数据
            data_files = list(self.data_path.glob("*.csv"))
            if not data_files:
                raise ValueError("未找到薪酬数据文件")
                
            df = pd.concat([pd.read_csv(f) for f in data_files])
            
            # 特征工程
            X = pd.get_dummies(df[self.features])
            y = df['salary']
            
            # 训练测试分割
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=0.2, random_state=42
            )
            
            # 训练模型
            self.model = GradientBoostingRegressor(
                n_estimators=100,
                learning_rate=0.1,
                max_depth=3,
                random_state=42
            )
            self.model.fit(X_train, y_train)
            
            # 评估模型
            score = self.model.score(X_test, y_test)
            logger.info(f"模型训练完成，R²分数: {score:.3f}")
            
            # 保存模型
            with open(self.model_path, 'wb') as f:
                pickle.dump(self.model, f)
                
            return score
            
        except Exception as e:
            logger.error(f"模型训练失败: {str(e)}", exc_info=True)
            raise
    
    def predict_salary(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """预测薪酬"""
        if not self.model:
            self._load_model()
            
        try:
            # 准备输入数据
            input_df = pd.DataFrame([input_data])
            features = pd.get_dummies(input_df[self.features])
            
            # 确保特征列一致
            train_columns = list(pd.read_csv(next(self.data_path.glob("*.csv")))[self.features].columns)
            for col in train_columns:
                if col not in features:
                    features[col] = 0
            
            # 预测
            prediction = self.model.predict(features)[0]
            
            # 生成解释
            explanation = self._generate_explanation(input_data, prediction)
            
            return {
                "predicted_salary": round(prediction, 2),
                "currency": "CNY",
                "explanation": explanation,
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            logger.error(f"薪酬预测失败: {str(e)}", exc_info=True)
            raise
    
    def _load_model(self):
        """加载已训练的模型"""
        try:
            with open(self.model_path, 'rb') as f:
                self.model = pickle.load(f)
            logger.info("已加载预训练模型")
        except Exception as e:
            logger.error(f"模型加载失败: {str(e)}")
            raise
    
    def _generate_explanation(self, input_data: Dict[str, Any], salary: float) -> str:
        """生成薪酬解释"""
        factors = []
        
        # 经验因素
        exp = input_data['years_experience']
        if exp < 3:
            factors.append(f"初级水平({exp}年经验)")
        elif exp < 7:
            factors.append(f"中级水平({exp}年经验)")
        else:
            factors.append(f"资深水平({exp}年经验)")
        
        # 学历因素
        edu_map = {
            1: "高中及以下",
            2: "大专",
            3: "本科",
            4: "硕士",
            5: "博士"
        }
        factors.append(f"学历:{edu_map.get(input_data['education_level'], '未知')}")
        
        # 技能因素
        factors.append(f"掌握{input_data['skill_count']}项核心技能")
        
        # 行业因素
        factors.append(f"{input_data['industry']}行业")
        
        # 城市因素
        location_map = {
            1: "一线城市",
            2: "二线城市",
            3: "三线及以下城市"
        }
        factors.append(location_map.get(input_data['location_tier'], "未知地区"))
        
        return f"预测月薪: ¥{salary:,.2f} (主要考虑因素: {', '.join(factors)})"

if __name__ == "__main__":
    # 示例使用
    predictor = SalaryPredictor()
    
    # 训练模型 (首次运行需要)
    # predictor.train_model()
    
    # 示例预测
    sample_input = {
        "years_experience": 4,
        "education_level": 3,  # 本科
        "skill_count": 6,
        "industry": "IT",
        "location_tier": 1  # 一线城市
    }
    
    result = predictor.predict_salary(sample_input)
    print(json.dumps(result, indent=2))